package com.handa.webShell.websocket.handle;


import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.handa.webShell.entity.DTO.FileInfo;
import com.handa.webShell.entity.DTO.SFTPFileInfoDTO;
import com.handa.webShell.entity.DTO.ServerConnectInfo;
import com.handa.webShell.entity.DTO.WebSocketDataBody;
import com.handa.webShell.entity.VO.SFTPWebSocketMessage;
import com.jcraft.jsch.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import javax.annotation.PostConstruct;
import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Component
@Slf4j
public class WebSFTPHandler extends TextWebSocketHandler {

    private static final Map<WebSocketSession, ChannelSftp> channelSftpMap = new ConcurrentHashMap<>();
    private static final Map<WebSocketSession, Session> sessionMap = new ConcurrentHashMap<>();
    // 创建SimpleDateFormat对象
    SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss:SSS");
    Map<String, String> map = new HashMap<>();

    /**
     * 将SftpATTRS的时间戳转换为格式化字符串
     *
     * @param attrs SftpATTRS对象
     * @return 格式化后的时间字符串
     */
    public static String formatSftpTime(SftpATTRS attrs) {
        try {
            // 获取文件的修改时间（秒级时间戳）
            long timeInSeconds = (long) attrs.getMTime();

            // 转换为毫秒
            long timeInMillis = timeInSeconds * 1000;

            // 创建SimpleDateFormat对象
            SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss:SSS");

            // 设置时区（可选，根据需求调整）
            sdf.setTimeZone(TimeZone.getTimeZone("Asia/Shanghai"));

            // 格式化时间
            return sdf.format(new Date(timeInMillis));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @PostConstruct
    public void init() {
        // 设置时区（可选，根据需求调整）
        sdf.setTimeZone(TimeZone.getTimeZone("Asia/Shanghai"));
        map.put("type", "pong");
    }

    @Override
    public void handleTextMessage(@NotNull WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        if (payload.isBlank()) {
            log.info("消息为空，忽略");
            return;
        } else {
            log.debug(payload);
        }

        WebSocketDataBody dataBody = JSONObject.parseObject(payload, WebSocketDataBody.class);
        String msgType = dataBody.getType();

        if (StringUtils.isBlank(msgType)) {
            log.info("消息类型为空，忽略");
            return;
        }

        if ("ping".equals(msgType)) {
            sendMessage(session, JSONObject.toJSONString(map));
            return;
        }

        Object data = dataBody.getData();
        if (Objects.isNull(data)) {
            log.info("消息数据体为空，忽略");
            return;
        }

        switch (msgType) {
            case "connectSftp":
                ServerConnectInfo serverConnectInfo = JSON.parseObject(JSONObject.toJSONString(data), ServerConnectInfo.class);
                connectToSFTP(session, serverConnectInfo);
                break;
            case "sftpPath":
                SFTPFileInfoDTO fileInfoDTO = JSON.parseObject(JSONObject.toJSONString(data), SFTPFileInfoDTO.class);
                listSFTPList(session, fileInfoDTO.getPath());
                break;
            case "sftpDownload":
                SFTPFileInfoDTO fileInfoDTO1 = JSON.parseObject(JSONObject.toJSONString(data), SFTPFileInfoDTO.class);
                handleDownloadSFTPBinary(session, fileInfoDTO1.getPath(), fileInfoDTO1.getFileName());
                break;
            default:
                log.error("消息类型 {} 没有对应的处理方法", msgType);
        }
    }

    // 二进制方式下载
    private void handleDownloadSFTPBinary(@NotNull WebSocketSession session, String path, String fileName) {
        try {
            ChannelSftp channelSftp = channelSftpMap.get(session);
            if (channelSftp == null) {
                SFTPWebSocketMessage sftpWebSocketMessage = new SFTPWebSocketMessage();
                sftpWebSocketMessage.setMsgType("sftpDownloadError");
                sftpWebSocketMessage.setFilePath(path);
                sftpWebSocketMessage.setFileName(fileName);
                sftpWebSocketMessage.setMessage("SFTP未连接或会话失效");
                session.sendMessage(new TextMessage(JSONObject.toJSONString(sftpWebSocketMessage)));

                return;
            }

            String remotePath = path.endsWith("/") ? path + fileName : path + "/" + fileName;

            // 获取总大小
            SftpATTRS attrs = channelSftp.stat(remotePath);
            long totalBytes = attrs.getSize();

            // 先发start文本，前端可据此初始化下载状态（包含总大小与目标分片大小）
            SFTPWebSocketMessage sftpWebSocketMessage = new SFTPWebSocketMessage();
            sftpWebSocketMessage.setMsgType("sftpDownloadStart");
            sftpWebSocketMessage.setFilePath(path);
            sftpWebSocketMessage.setFileName(fileName);

            Map<String, Object> startData = new HashMap<>();
            final int targetChunk = 2 * 1024 * 1024; // 目标聚合分片：2MB，可按需调整
            startData.put("totalBytes", totalBytes);
            startData.put("chunkBytes", targetChunk);
            sftpWebSocketMessage.setData(startData);

            String sendStr = JSONObject.toJSONString(sftpWebSocketMessage);
            session.sendMessage(new TextMessage(sendStr));
            log.info("sftpDownloadStart: " + sendStr);

            // 分片发送二进制内容（聚合多次read形成较大分片）
            try (InputStream is = channelSftp.get(remotePath)) {
                byte[] smallBuf = new byte[64 * 1024]; // 单次读取缓存，SFTP每包约32KB
                java.io.ByteArrayOutputStream chunk = new java.io.ByteArrayOutputStream(targetChunk);
                BufferedInputStream bis = new BufferedInputStream(is, 64 * 1024); // 64KB 缓冲区

                int read;
                long sentBytes = 0;
                while ((read = bis.read(smallBuf)) != -1) {
                    if (read <= 0) {
                        continue;
                    }
                    int offset = 0;
                    while (offset < read) {
                        int remaining = targetChunk - chunk.size();
                        int toWrite = Math.min(remaining, read - offset);
                        chunk.write(smallBuf, offset, toWrite);
                        offset += toWrite;

                        if (chunk.size() == targetChunk) {
                            byte[] payload = chunk.toByteArray();
                            session.sendMessage(new org.springframework.web.socket.BinaryMessage(java.nio.ByteBuffer.wrap(payload)));
                            sentBytes += payload.length;
                            // 进度消息（文本）
                            SFTPWebSocketMessage progress = new SFTPWebSocketMessage();
                            progress.setMsgType("sftpDownloadProgress");
                            progress.setFilePath(path);
                            progress.setFileName(fileName);
                            Map<String, Object> prog = new HashMap<>();
                            prog.put("sentBytes", sentBytes);
                            prog.put("totalBytes", totalBytes);
                            prog.put("percent", totalBytes == 0 ? 100.0 : (sentBytes * 100.0 / totalBytes));
                            log.info(JSONObject.toJSONString("正常分片数据： " + prog));
                            progress.setData(prog);
                            session.sendMessage(new TextMessage(JSONObject.toJSONString(progress)));
                            chunk.reset();
                        }
                    }
                }
                // 发送最后不足一个聚合分片的残余（如果有）
                if (chunk.size() > 0) {
                    byte[] payload = chunk.toByteArray();
                    // 发送尾部分片数据
                    session.sendMessage(new org.springframework.web.socket.BinaryMessage(java.nio.ByteBuffer.wrap(payload)));
                    sentBytes += payload.length;

                    // 发送进度更新
                    SFTPWebSocketMessage progress = new SFTPWebSocketMessage();
                    progress.setMsgType("sftpDownloadProgress");
                    progress.setFilePath(path);
                    progress.setFileName(fileName);
                    Map<String, Object> prog = new HashMap<>();
                    prog.put("sentBytes", sentBytes);
                    prog.put("totalBytes", totalBytes);
                    prog.put("percent", totalBytes == 0 ? 100.0 : (sentBytes * 100.0 / totalBytes));
                    log.info(JSONObject.toJSONString("最后一小点数据： " + prog));
                    progress.setData(prog);
                    session.sendMessage(new TextMessage(JSONObject.toJSONString(progress)));

                    // 关键：等待一小段时间，确保二进制数据完全发送
                    Thread.sleep(100);
                }
            }

            // 确保所有二进制数据发送完成后再发送结束信号
            SFTPWebSocketMessage sftpWebSocketMessageEnd = new SFTPWebSocketMessage();
            sftpWebSocketMessageEnd.setMsgType("sftpDownloadEnd");
            sftpWebSocketMessageEnd.setFilePath(path);
            sftpWebSocketMessageEnd.setFileName(fileName);
            Map<String, Object> endData = new HashMap<>();
            endData.put("totalBytes", totalBytes);
            endData.put("status", "ok");
            sftpWebSocketMessageEnd.setData(endData);
            String endStr = JSONObject.toJSONString(sftpWebSocketMessageEnd);
            session.sendMessage(new TextMessage(endStr));
            log.info("sftpDownloadEnd: " + endStr);
        } catch (Exception e) {
            try {
                SFTPWebSocketMessage sftpWebSocketMessage = new SFTPWebSocketMessage();
                sftpWebSocketMessage.setMsgType("sftpDownloadError");
                sftpWebSocketMessage.setFilePath(path);
                sftpWebSocketMessage.setFileName(fileName);
                sftpWebSocketMessage.setMessage(e.getMessage());

                session.sendMessage(new TextMessage(JSONObject.toJSONString(sftpWebSocketMessage)));
            } catch (IOException ex) {
                throw new RuntimeException(ex);
            }
        }
    }

    // 连接到 SFTP
    private void connectToSFTP(WebSocketSession session, ServerConnectInfo serverConnectInfo) {
        try {
            Session sshSession = connectToSession(serverConnectInfo);

            sshSession.setPassword(serverConnectInfo.getPassword());
            Properties configProps = new Properties();
            configProps.put("StrictHostKeyChecking", "no");
            configProps.put("max_input_buffer_size", "65536"); // 64KB
            configProps.put("max_packet_size", "65536"); // 64KB
            sshSession.setConfig(configProps);
            sshSession.connect();

            Channel channel = sshSession.openChannel("sftp");
            channel.connect();
            ChannelSftp channelSftp = (ChannelSftp) channel;

            sessionMap.put(session, sshSession);
            channelSftpMap.put(session, channelSftp);

            listSFTPList(session, "/");

        } catch (JSchException e) {
            try {
                session.sendMessage(new TextMessage("SFTP连接失败: " + e.getMessage()));
                session.close();
            } catch (IOException ex) {
                throw new RuntimeException(ex);
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // 接收路径， 发送路径下的文件列表信息
    public void listSFTPList(WebSocketSession session, String path) throws SftpException, IOException {
        ChannelSftp channelSftp = channelSftpMap.get(session);

        List<FileInfo> fileInfos = new ArrayList<>();
        Vector<ChannelSftp.LsEntry> files = channelSftp.ls(path);

        for (ChannelSftp.LsEntry entry : files) {
            if (entry.getFilename().equals(".") || entry.getFilename().equals("..")) {
                continue;
            }
            FileInfo fileInfo = new FileInfo();
            fileInfo.setName(entry.getFilename());
            fileInfo.setFolderFlag(entry.getAttrs().isDir());
            fileInfo.setSize(entry.getAttrs().getSize());
            fileInfo.setLastModified(formatSftpTime(entry.getAttrs()));
            fileInfos.add(fileInfo);
        }

        SFTPWebSocketMessage sftpWebSocketMessage = new SFTPWebSocketMessage();
        sftpWebSocketMessage.setMsgType("fileList");
        sftpWebSocketMessage.setData(fileInfos);

        sendMessage(session, JSONObject.toJSONString(sftpWebSocketMessage));

    }

    // 向 WebSocket 发送消息
    private void sendMessage(WebSocketSession session, String message) throws IOException {
        session.sendMessage(new TextMessage(message));
    }

    // 关闭 WebShell SSH
    private void closeSSH(WebSocketSession session) {
        Session sshSession = sessionMap.get(session);
        ChannelSftp sftpChannel = channelSftpMap.get(session);

        if (sftpChannel != null) {
            sftpChannel.disconnect();
            channelSftpMap.remove(session);
        }

        if (sshSession != null) {
            sshSession.disconnect();
            sessionMap.remove(session);
        }
    }

    @Override
    public void afterConnectionClosed(@NotNull WebSocketSession session, @NotNull CloseStatus status) throws Exception {
        log.info("sftp session close: {}", session.getId());
        closeSSH(session);
        super.afterConnectionClosed(session, status);
    }

    // 获取jSCh Session
    private Session connectToSession(ServerConnectInfo serverConnectInfo) throws JSchException {
        JSch jSch = new JSch();

        return jSch.getSession(serverConnectInfo.getUsername()
                , serverConnectInfo.getHost(),
                serverConnectInfo.getPort());
    }
}